from mcp.server.fastmcp import FastMCP
from typing import List, Dict
from scipy.stats import chi2_contingency
import numpy as np

# Create an MCP server for Chi-Square tests
mcp = FastMCP("Chi-Square Test Server")

@mcp.tool()
def chi_square_goodness_of_fit(observed: List[float], expected: List[float]) -> Dict:
    """Perform Chi-Square Goodness of Fit test"""
    chi2, p = chi2_contingency(np.array([observed, expected]))[0:2]
    dof = len(observed) - 1
    return {
        "chi_square": chi2,
        "p_value": p,
        "dof": dof,
        "significance": "significant" if p < 0.05 else "not significant"
    }

@mcp.tool()
def chi_square_independence(contingency_table: List[List[float]]) -> Dict:
    """Perform Chi-Square Test of Independence"""
    chi2, p, dof, _ = chi2_contingency(np.array(contingency_table))
    return {
        "chi_square": chi2,
        "p_value": p,
        "dof": dof,
        "significance": "significant" if p < 0.05 else "not significant"
    }

@mcp.tool()
def chi_square_homogeneity(groups: List[List[float]]) -> Dict:
    """Perform Chi-Square Test of Homogeneity"""
    chi2, p, dof, _ = chi2_contingency(np.array(groups))
    return {
        "chi_square": chi2,
        "p_value": p,
        "dof": dof,
        "significance": "significant" if p < 0.05 else "not significant"
    }

def get_input_list(prompt):
    print(prompt)
    while True:
        try:
            values = input("请输入数值，用空格分隔: ").split()
            return [float(x) for x in values]
        except ValueError:
            print("输入无效，请重新输入数字！")

if __name__ == "__main__":
    print("欢迎使用卡方检验MCP服务器")
    print("1. 拟合优度检验\n2. 独立性检验\n3. 同质性检验")
    
    while True:
        choice = input("请选择检验类型(1-3)，或输入q退出: ")
        if choice.lower() == 'q':
            break
            
        try:
            choice = int(choice)
            if choice == 1:
                observed = get_input_list("\n拟合优度检验 - 请输入观察值:")
                expected = get_input_list("请输入期望值:")
                result = chi_square_goodness_of_fit(observed, expected)
            elif choice == 2:
                print("\n独立性检验 - 请输入列联表(每行用分号分隔，数值用空格分隔):")
                table = []
                while True:
                    row = input("输入一行(或输入d完成): ").strip()
                    if row.lower() == 'd':
                        if len(table) >= 2:
                            break
                        print("至少需要2行数据")
                    else:
                        try:
                            values = [float(x) for x in row.split()]
                            table.append(values)
                        except ValueError:
                            print("输入无效，请重新输入")
                result = chi_square_independence(table)
            elif choice == 3:
                print("\n同质性检验 - 请输入各组数据(每组用分号分隔，数值用空格分隔):")
                groups = []
                while True:
                    group = input("输入一组(或输入d完成): ").strip()
                    if group.lower() == 'd':
                        if len(groups) >= 2:
                            break
                        print("至少需要2组数据")
                    else:
                        try:
                            values = [float(x) for x in group.split()]
                            groups.append(values)
                        except ValueError:
                            print("输入无效，请重新输入")
                result = chi_square_homogeneity(groups)
            else:
                print("无效选择，请输入1-3")
                continue
                
            print("\n检验结果:")
            for k, v in result.items():
                print(f"{k}: {v}")
            print()
                
        except ValueError:
            print("无效输入，请重新选择")
    
    # 保留原有的stdio模式，以便其他程序调用
    mcp.run(transport="stdio")